import yaml
import pandas as pd
import numpy as np
from pymongo import MongoClient

# 用于评估

DB_NAME = 'stock'


def load_config():
    with open("config.yaml", 'r', encoding='utf-8') as f:
        data = yaml.load(f.read(), yaml.FullLoader)
    print("loading_config ... >>> 读取本地配置")
    return data


def loc_collection(db_name, collection_name, cfg={}):
    if len(cfg) == 0:
        cfg = load_config()
    conn = MongoClient(cfg['mongodb'])
    coll_ = conn[db_name][collection_name]
    return coll_


def loc_db(db_name, cfg={}):
    if len(cfg) == 0:
        cfg = load_config()
    return MongoClient(cfg['mongodb'])[db_name]


def transfer_data(data):
    dic = {}
    ts_codes = [s['y_name'] for s in data]
    for ts_code in ts_codes:
        filter_data = [s for s in data if s['y_name'] == ts_code]
        thd = 0.7
        y_pred_data = [s['y_pred'][-9:] for s in filter_data if float(s['acc']) > thd]
        acc_li = [float(s['acc']) for s in filter_data]
        max_acc = np.max(acc_li) - 0.005
        max_pred_data = [s['y_pred'][-9:] for s in filter_data if float(s['acc']) > max_acc]
        while len(y_pred_data) < 1:
            thd -= 0.5
            y_pred_data = [s['y_pred'][-9:] for s in filter_data if float(s['acc']) > thd]
        tt = filter_data[0]['t'][-9:]
        tt = [str(s)[:10] for s in tt]
        # df_pred = pd.DataFrame(data=y_pred_data, columns=tt).T
        df_max = pd.DataFrame(data=max_pred_data, columns=tt).T
        dic[ts_code] = df_max.mean(axis=1)
    return dic


cfg = load_config()
db = loc_db(DB_NAME, cfg)
coll_names = db.list_collection_names()

return_dic = {}
for coll_name in coll_names:
    coll = loc_collection(DB_NAME, coll_name, cfg)
    find_data = list(coll.find({}))
    return_dic[coll_name] = transfer_data(find_data)
